package cn.edu.hit.mongodb;

import cn.edu.hit.core.Document;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.Filters;
import org.bson.BsonDocument;
import org.bson.types.ObjectId;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

public class Utils {
    public static String persist(MongoCollection<Document> collection, Map docMap) {
        Document document = new Document(docMap);
        if (org.apache.commons.lang3.StringUtils.isNotEmpty(document.getString("_id"))) {
            ObjectId _id = new ObjectId(document.getString("_id"));
            document.put("_id", _id);
            Document result = collection.findOneAndReplace(
                    Filters.eq("_id", _id), document);
            if (result != null) {
                return _id.toHexString();
            }
        }

        ObjectId _id = ObjectId.get();
        document.put("_id", _id);
        collection.insertOne(document);
        return _id.toHexString();
    }

    public static List<String> persistMany(MongoCollection<Document> collection, List<Map> docMaps) {
        return persistManyR(collection, docMaps, (_id, docMap) -> _id);
    }

    public static List<Map> persistMany2(MongoCollection<Document> collection, List<Map> docMaps) {
        return persistManyR(collection, docMaps, (_id, docMap) -> {
            docMap.put("_id", _id);
            return docMap;
        });
    }

    public static <R> List<R> persistManyR(MongoCollection<Document> collection, List<Map> docMaps, BiFunction<String, Map, R> transferFunction) {
        return docMaps.stream().map(docMap -> {
            String id = Utils.persist(collection, docMap);
            return transferFunction.apply(id, docMap);
        }).collect(Collectors.toList());
    }

    public static List<Document> getLastRecords(MongoDatabase database, String collectionName, BsonDocument filter, long timestamp) {
        MongoCollection<Document> conclusions = database.getCollection(collectionName, Document.class);
        if (filter == null) {
            filter = new BsonDocument();
        }

        Document result = conclusions.aggregate(MongoDBConnection.asList(
                new Document("$match", filter),
                new Document("$addFields",
                        new Document("subVal",
                                new Document("$abs",
                                        new Document("$subtract",
                                                Arrays.asList("$timestamp", timestamp))))),
                new Document("$sort", new Document("subVal", 1))
        )).first();
        if (result == null) {
            return Collections.emptyList();
        }
        long stamp = result.getLong("timestamp");

        return StreamSupport.stream(conclusions.find(Filters.and(filter, Filters.eq("timestamp", stamp))).spliterator(), true)
                .peek(doc -> doc.put("_id", doc.getObjectId("_id")))
                .collect(Collectors.toList());
    }
}
